# vector_processor.py
import os
import jieba
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from langchain_community.document_loaders import Docx2txtLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from typing import List, Dict
from config import EMBEDDING_MODEL_PATH


class KnowledgeBaseProcessor:
    """处理文档加载、向量化、关键词提取和检索的类。"""

    def __init__(self, knowledge_files: list, vector_store_path: str):
        self.knowledge_files = knowledge_files
        self.vector_store = None
        self.documents = []
        self.keywords = None
        self.vector_store_path = vector_store_path
        # 停用词库，确保和原始代码一致
        self.stopwords = {
            '的', '地', '得', '了', '着', '是', '有', '也', '和', '而', '之', '以', '及', '于', '在',
            '被', '把', '但', '或', '则', '又', '更', '都', '对', '从', '为', '个', '种', '等', '、',
            '。', '《', '》', '，', '“', '”', '！', '？', '：', '；', '（', '）', '·', '—', '…', '·',
            ' ', '\n', '\t'
        }
        self.tfidf_vectorizer = TfidfVectorizer(tokenizer=self._jieba_tokenizer, stop_words=self.stopwords)

    def _jieba_tokenizer(self, text):
        return [word for word in jieba.cut(text) if word.strip()]

    def get_files_metadata(self) -> List[Dict]:
        """新增方法：获取所有知识文件的元数据（修改时间）。"""
        metadata = []
        for file_path in self.knowledge_files:
            try:
                # 使用 os.path.getmtime 获取文件修改时间戳
                mtime = os.path.getmtime(file_path)
                metadata.append({"file_path": file_path, "mtime": mtime})
            except Exception as e:
                # 如果文件不存在，将元数据标记为None
                print(f"无法获取文件元数据: {file_path}, 错误: {e}")
                metadata.append({"file_path": file_path, "mtime": None})
        return metadata

    def load_documents(self) -> None:
        """加载知识文件并切分文档。"""
        self.documents = []
        for file in self.knowledge_files:
            try:
                if file.endswith('.docx'):
                    loader = Docx2txtLoader(file)
                elif file.endswith('.txt'):
                    loader = TextLoader(file, encoding='utf-8')
                else:
                    print(f"Warning: Unsupported file type for {file}")
                    continue

                documents = loader.load()
                text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
                self.documents.extend(text_splitter.split_documents(documents))
            except Exception as e:
                print(f"Failed to load document {file}: {e}")

        print(f"Loaded {len(self.documents)} document chunks.")

    def create_vector_store(self) -> None:
        """创建或加载向量存储。"""
        embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_PATH)
        if os.path.exists(self.vector_store_path):
            self.vector_store = FAISS.load_local(self.vector_store_path, embeddings,
                                                 allow_dangerous_deserialization=True)
            print("Successfully loaded existing vector store.")
        else:
            if not self.documents:
                raise ValueError("No documents to create vector store from.")
            self.vector_store = FAISS.from_documents(self.documents, embeddings)
            self.vector_store.save_local(self.vector_store_path)
            print("Successfully created and saved new vector store.")

    def extract_keywords(self) -> None:
        """使用TF-IDF提取文档关键词。"""
        if not self.documents:
            print("No documents loaded, skipping keyword extraction.")
            return

        texts = [doc.page_content for doc in self.documents]
        tfidf_matrix = self.tfidf_vectorizer.fit_transform(texts)
        feature_names = self.tfidf_vectorizer.get_feature_names_out()
        tfidf_scores = np.asarray(tfidf_matrix.mean(axis=0)).ravel()
        sorted_indices = tfidf_scores.argsort()[-50:][::-1]
        self.keywords = [feature_names[i] for i in sorted_indices]
        print("Keywords extracted successfully.")

    def retrieve_relevant_docs(self, query: str, top_k: int = 3) -> List:
        """从向量库中检索与问题最相关的文档。"""
        if self.vector_store is None:
            print("Vector store not initialized.")
            return []

        return self.vector_store.similarity_search(query, k=top_k)